fix(indexing): read a Boolean array index the way NumPy does - #443

🤖 AI text below 🤖

Indexing an array with a Boolean NumPy array now selects what NumPy selects. That covers a mask over the leading axes (for example a (2, 3) mask on a (2, 3, 4) array), an empty mask whose length is 0 along an axis, and a full-shape mask on an array with a zero-length axis. Orthogonal selection also accepts an empty Boolean array and no longer fails on a zero-length dimension. A Boolean array with more than one dimension inside a selection tuple, which is not supported, now raises an error that says so instead of one about orthogonal selection.

What changed

expand_bool_index in src/zarr/core/indexing.py rewrites a Boolean array used as the whole index into a selection that the existing indexers already handle, following NumPy's reading of a mask m as m.nonzero():

• A mask with the array's full shape stays a mask selection. This includes empty masks on arrays with a zero-length axis, which work on main.
• A 1-D mask over the first of several axes becomes an orthogonal selection.
• Any other mask, including an empty one, becomes a coordinate selection that broadcasts against the remaining axes.
• A mask that does not match the axes it covers raises IndexError. As in NumPy, an empty mask may have length 0 along an axis; any other mismatch raises.

Array.__getitem__, Array.__setitem__ and get_indexer call it before dispatch. BoolArrayDimIndexer accepts a zero-length Boolean array and computes its item count without indexing an empty cumulative sum, which crashed on a zero-length dimension on main. The vindex contract (coordinate arrays, or one full-shape mask) is unchanged.

Verification

• One parametrized test compares reads and writes against NumPy, and one test covers the wrong-shape error; 18 of the 27 cases fail on main.
• A differential fuzz of 800 random cases against NumPy 2.5.4 (up to 3 dimensions, zero-length axes, random chunks, sharded and unsharded, reads plus scalar and array writes) found about 280 reads that fail on main and pass here, no regressions, and matching shapes and values wherever both return data.
• The full test suite and mypy pass locally.

Known limits

• A mask over two or more leading axes of a higher-dimensional array builds one coordinate per selected element, about 22 times the result's size in one measurement. The full-shape mask path costs the same.
• Writes match NumPy for a scalar value or a value with the selection's full shape. zarr does not broadcast a smaller value on any selection path, so z[m2] = np.arange(4) on a (2, 3, 4) array still raises.
• A mask with two or more dimensions inside a tuple, such as z[m2, 1] or z[m2, :], still raises, now with an error that says such a mask is only supported as the whole index (it used to talk about orthogonal selection). Supporting it would mean turning every slice in the tuple into coordinates, so it is left for a move onto zarr-indexing, whose vindex keeps those axes lazy.
• A 0-d Boolean index (z[np.True_]) still raises; NumPy adds a length-1 axis.
• Several Boolean arrays in one tuple are still read orthogonally, as on main, so z[E, E] with two empty masks gives (0, 0) where NumPy gives (0,).
• a[m, ...] with an index array before an Ellipsis still fails; that is fixed separately.
• Field selection on structured arrays is broken on main for every selection type and is not touched here.

🤖 Generated with Claude Code